#include "ir.h"

#include <iostream>
#include <ostream>
#include <unordered_set>
#include <vector>

#include "effect.h"
#include "irblock.h"
#include "irfunction.h"
#include "iroperand.h"
#include "utils.h"

#define VAR BLUE
#define OP DEF
#define LABEL COLOR(5)

using namespace Utils;

namespace IR {
    void Operation::Print() const {
        if (dest) dest->Print();
        std::cout << OP " := ";
        if (lhs) lhs->Print();
        std::cout << OP " " << (char)op << " ";
        if (rhs) rhs->Print();
    }
    std::vector<Variable*> Operation::UsedVar() const {
        std::vector<Variable*> ret;
        if (auto v = lhs->As<Variable>()) ret.push_back(v);
        if (auto v = rhs->As<Variable>()) ret.push_back(v);
        return ret;
    }
    std::vector<Variable*> Operation::ChangedVar() const {
        std::vector<Variable*> ret;
        if (dest) ret.push_back(dest);
        return ret;
    }
    SideEffect Operation::Effects() const {
        SideEffect ret;
        if (dest) ret.AddVar(dest);
        return ret;
    }

    void Alloc::Print() const {
        if (dest) dest->Print();
        std::cout << OP " := new ";
        std::cout << VAR << size;
    }
    std::vector<Variable*> Alloc::UsedVar() const { return {}; }
    std::vector<Variable*> Alloc::ChangedVar() const { return {dest}; }
    SideEffect Alloc::Effects() const {
        SideEffect ret;
        ret.AddVar(dest);
        return ret;
    }

    void Load::Print() const {
        if (dest) dest->Print();
        std::cout << OP " := ";
        if (base) base->Print();
        std::cout << OP "[";
        if (offset) offset->Print();
        std::cout << OP "]";
    }
    std::vector<Variable*> Load::UsedVar() const {
        std::vector<Variable*> ret{base};
        if (auto v = offset->As<Variable>()) ret.push_back(v);
        return ret;
    }
    std::vector<Variable*> Load::ChangedVar() const { return {dest}; }
    SideEffect Load::Effects() const {
        SideEffect ret;
        ret.AddVar(dest);
        return ret;
    }

    void Store::Print() const {
        if (base) base->Print();
        std::cout << OP "[";
        if (offset) offset->Print();
        std::cout << OP "]";
        std::cout << OP " := ";
        if (src) src->Print();
    }
    std::vector<Variable*> Store::UsedVar() const {
        std::vector<Variable*> ret{base};
        if (auto v = src->As<Variable>()) ret.push_back(v);
        if (auto v = offset->As<Variable>()) ret.push_back(v);
        return ret;
    }
    std::vector<Variable*> Store::ChangedVar() const { return {}; }
    SideEffect Store::Effects() const {
        SideEffect ret;
        ret.AddMem(base);
        return ret;
    }

    void Phi::Print() const {
        if (dest) dest->Print();
        std::cout << OP " := phi(";
        int i = 0;
        for (auto [block, src] : srcs) {
            if (i++) std::cout << OP ", ";
            std::cout << OP "[" LABEL << block->name << OP ", ";
            src->Print();
            std::cout << OP "]";
        }
        std::cout << OP ")";
    }
    std::vector<Variable*> Phi::UsedVar() const {
        std::vector<Variable*> ret;
        for (auto [block, src] : srcs) ret.push_back(src);
        return ret;
    }
    std::vector<Variable*> Phi::ChangedVar() const { return {dest}; }
    SideEffect Phi::Effects() const {
        SideEffect ret;
        ret.AddVar(dest);
        return ret;
    }

    void Call::Print() const {
        if (dest) dest->Print();
        std::cout << OP " := " << func->name << "(";
        for (int i = 0; i < args.size(); ++i) {
            if (i) std::cout << OP ", ";
            args[i]->Print();
        }
        std::cout << OP ")";
    }
    std::vector<Variable*> Call::UsedVar() const {
        std::vector<Variable*> ret;
        for (auto x : args)
            if (auto v = x->As<Variable>()) ret.push_back(v);
        return ret;
    }
    std::vector<Variable*> Call::ChangedVar() const {
        std::vector<Variable*> ret;
        if (dest) ret.push_back(dest);
        return ret;
    }
    SideEffect Call::Effects() const {
        SideEffect ret;
        if (func->IsExternal()) ret.io = true;
        if (dest) ret.AddVar(dest);
        ret.Merge(func->globalEffects);
        for (auto id : func->argMemEffects) {
            auto arg = args[id]->As<Variable>();
            ret.AddMem(arg);
        }
        return ret;
    }

    void Jump::Print() const {
        std::cout << OP "jmp ";
        std::cout << LABEL << nxt->name;
    }
    std::vector<Variable*> Jump::UsedVar() const { return {}; }
    std::vector<Variable*> Jump::ChangedVar() const { return {}; }
    SideEffect Jump::Effects() const { return {}; }
    std::vector<IRBlock*> Jump::GetNext() const { return {nxt}; }

    void Return::Print() const {
        std::cout << OP "return ";
        if (ret) ret->Print();
    }
    std::vector<Variable*> Return::UsedVar() const {
        std::vector<Variable*> ret;
        if (auto v = this->ret->As<Variable>()) ret.push_back(v);
        return ret;
    }
    std::vector<Variable*> Return::ChangedVar() const { return {}; }
    SideEffect Return::Effects() const { return {}; }
    std::vector<IRBlock*> Return::GetNext() const { return {}; }

    void Branch::Print() const {
        std::cout << OP "branch ";
        if (lhs) lhs->Print();
        std::cout << OP " " << (char)op << " ";
        if (rhs) rhs->Print();
        std::cout << LABEL " " << t->name << " " << f->name;
    }
    std::vector<Variable*> Branch::UsedVar() const {
        std::vector<Variable*> ret;
        if (auto v = lhs->As<Variable>()) ret.push_back(v);
        if (auto v = rhs->As<Variable>()) ret.push_back(v);
        return ret;
    }
    std::vector<Variable*> Branch::ChangedVar() const { return {}; }
    SideEffect Branch::Effects() const { return {}; }
    std::vector<IRBlock*> Branch::GetNext() const { return {t, f}; }

}  // namespace IR